package testedge

import (
    "fmt"
    "sync"
    "time"
)

// Task 定义任务接口
type Task func()

// Pool 协程池结构
type Pool struct {
    Workers  int               // 协程池尺寸
    taskQueue   chan Task         // 任务队列
    activeCount int               // 当前活跃协程数量
    wg          sync.WaitGroup    // 等待组，用于同步协程
    stop        chan struct{}     // 用于通知协程停止
    isClosed    bool              // 标记协程池是否已关闭
}

// New 创建一个新的协程池
func New(workers int) *Pool {
    return &Pool{
        Workers:  workers,
        taskQueue:   make(chan Task),
        stop:        make(chan struct{}),
        activeCount: 0,
    }
}

// Run 提交任务到协程池
func (p *Pool) Run(task Task) {
    if p.isClosed {
        fmt.Println("协程池已关闭，无法提交新任务")
        return
    }
    p.taskQueue <- task
}

// Start 启动协程池
func (p *Pool) Start() {
    // 初始化最小数量的工作协程
    for i := 0; i < p.Workers; i++ {
        p.wg.Add(1)
        go p.worker()
    }

    // 启动监控协程
    go p.monitor()
}

// worker 工作协程
func (p *Pool) worker() {
    defer p.wg.Done()

    for {
        select {
        case task, ok := <-p.taskQueue:
            if !ok { // 通道关闭，退出协程
                p.activeCount--
                return
            }
            p.activeCount++
            task()
            p.activeCount--
        case <-p.stop:
            p.activeCount--
            return
        }
    }
}

// monitor 监控协程
func (p *Pool) monitor() {
    ticker := time.NewTicker(1 * time.Second) // 每秒检查一次
    defer ticker.Stop()

    for {
        select {
        case <-ticker.C:
            // 如果当前活跃协程数量低于最小值，则拉起新的协程
            if p.activeCount < p.Workers {
                p.wg.Add(1)
                go p.worker()
            }
        case <-p.stop:
            return
        }
    }
}

// Shutdown 等待所有任务完成并关闭协程池
func (p *Pool) Shutdown() {
    if p.isClosed {
        fmt.Println("协程池已关闭，无需重复关闭")
        return
    }

    close(p.taskQueue) // 关闭任务队列
    p.isClosed = true  // 标记协程池已关闭

    p.wg.Wait()        // 等待所有工作协程完成
    close(p.stop)      // 通知所有工作协程停止
    fmt.Println("协程池已关闭")
}

func main() {
    pool := New(3) // 创建一个最小协程数为3，最大协程数为5的协程池
    pool.Start()      // 启动协程池

    // 提交任务到协程池
    for i := 0; i < 10; i++ {
        task := func(i int) Task {
            return func() {
                fmt.Printf("任务 %d 正在执行\n", i)
                time.Sleep(time.Second)
            }
        }(i)
        pool.Run(task)
    }

    // 等待所有任务完成并关闭协程池
    pool.Shutdown()
}
